package com.knowledge.business.service.impl;

import com.alibaba.fastjson2.JSON;
import com.google.gson.JsonArray;
import com.google.gson.JsonObject;
import com.knowledge.business.domain.KnowledgeDocumentSegment;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.DataType;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.collection.request.AddFieldReq;
import io.milvus.v2.service.collection.request.CreateCollectionReq;
import io.milvus.v2.service.collection.request.HasCollectionReq;
import io.milvus.v2.service.collection.request.LoadCollectionReq;
import io.milvus.v2.service.index.request.CreateIndexReq;
import io.milvus.v2.service.vector.request.DeleteReq;
import io.milvus.v2.service.vector.request.InsertReq;
import io.milvus.v2.service.vector.request.SearchReq;
import io.milvus.v2.service.vector.request.data.FloatVec;
import io.milvus.v2.service.vector.response.InsertResp;
import io.milvus.v2.service.vector.response.SearchResp;
import jakarta.annotation.PostConstruct;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.*;
import java.util.concurrent.TimeUnit;

@Service
@Slf4j
public class VectorStorageService {
    private static final String COLLECTION_NAME = "knowledge_segments";
    private static final String ID_FIELD = "id";
    private static final String DOCUMENT_ID_FIELD = "document_id";
    private static final String SEGMENT_INDEX_FIELD = "segment_index";
    private static final String SEGMENT_ID_FIELD = "segment_id";
    private static final String VECTOR_FIELD = "vector";
    private static final String CONTENT_FIELD = "content";
    private static final String METADATA_FIELD = "metadata";
    private static final int VECTOR_DIMENSION = 1024; // 向量维度

    @Autowired
    private MilvusClientV2 milvusClient;


    @PostConstruct
    public void init() {
        try {
            log.info("初始化 Milvus service........................");
            createCollection();
            createVectorIndexIfNotExists();
            TimeUnit.SECONDS.sleep(5);
            loadCollection();
            log.info("Milvus service 初始化成功........................");
        } catch (Exception e) {
            log.error("Milvus service initialization failed: " + e.getMessage(), e);
        }
    }


    /**
     * 创建集合（包含完整的schema定义）
     */
    public void createCollection() {
        try {
            Boolean hasCollection = milvusClient.hasCollection(HasCollectionReq.builder()
                    .collectionName(COLLECTION_NAME)
                    .build());

            if (hasCollection) {
                log.info("集合 " + COLLECTION_NAME + " 已经存在................................");
                return;
            }

            // 创建schema并定义所有字段
            CreateCollectionReq.CollectionSchema schema = milvusClient.createSchema();

            // 主键字段
            schema.addField(AddFieldReq.builder()
                    .fieldName(ID_FIELD)
                    .dataType(DataType.VarChar)
                    .isPrimaryKey(true)
                    .autoID(false)
                    .maxLength(50)
                    .build());

            // 文档ID字段
            schema.addField(AddFieldReq.builder()
                    .fieldName(DOCUMENT_ID_FIELD)
                    .dataType(DataType.VarChar)
                    .maxLength(50)
                    .build());

            // 段落索引字段
            schema.addField(AddFieldReq.builder()
                    .fieldName(SEGMENT_INDEX_FIELD)
                    .dataType(DataType.Int64)
                    .maxLength(50)
                    .build());

            // 段落ID字段
            schema.addField(AddFieldReq.builder()
                    .fieldName(SEGMENT_ID_FIELD)
                    .dataType(DataType.VarChar)
                    .maxLength(50)
                    .build());

            // 向量字段（关键：明确指定维度）
            schema.addField(AddFieldReq.builder()
                    .fieldName(VECTOR_FIELD)
                    .dataType(DataType.FloatVector)
                    .dimension(VECTOR_DIMENSION) // 这里定义了向量维度
                    .build());

            // 内容字段
            schema.addField(AddFieldReq.builder()
                    .fieldName(CONTENT_FIELD)
                    .dataType(DataType.VarChar)
                    .maxLength(65535)
                    .build());

            // 元数据字段
            schema.addField(AddFieldReq.builder()
                    .fieldName(METADATA_FIELD)
                    .dataType(DataType.VarChar)
                    .maxLength(65535)
                    .build());

            // 创建集合
            milvusClient.createCollection(CreateCollectionReq.builder()
                    .collectionName(COLLECTION_NAME)
                    .collectionSchema(schema)
                    .build());

            log.info("Collection " + COLLECTION_NAME + " created successfully.");
        } catch (Exception e) {
            log.error("创建集合失败: " + e.getMessage(), e);
            throw new RuntimeException("创建集合失败: " + e.getMessage(), e);
        }
    }

    /**
     * 为向量字段创建索引（使用专门的createIndex方法）
     */
    public void createVectorIndexIfNotExists() {
        try {
            // 向量字段索引参数（不需要重复定义维度，因为schema中已定义）
            IndexParam vectorIndexParam = IndexParam.builder()
                    .fieldName(VECTOR_FIELD)
                    .indexType(IndexParam.IndexType.AUTOINDEX)
                    .metricType(IndexParam.MetricType.COSINE)
                    .build();

            List<IndexParam> indexParams = new ArrayList<>();
            indexParams.add(vectorIndexParam);

            // 使用专门的创建索引方法，而不是再次创建集合
            milvusClient.createIndex(CreateIndexReq.builder()
                    .collectionName(COLLECTION_NAME)
                    .indexParams(indexParams)
                    .build());
            log.info("向量字段 " + VECTOR_FIELD + " 索引创建成功");
        } catch (Exception e) {
            if (e.getMessage().contains("index already exists")) {
                log.info("向量字段 " + VECTOR_FIELD + " 索引已存在");
            } else {
                log.error("创建向量索引失败: " + e.getMessage(), e);
                throw new RuntimeException("创建向量索引失败: " + e.getMessage(), e);
            }
        }
    }


    /**
     * 加载集合到内存
     */
    public void loadCollection() {
        milvusClient.loadCollection(LoadCollectionReq.builder()
                .collectionName(COLLECTION_NAME)
                .build());

        log.info("Milvus service 向量集合加载成功..........................................");
    }

    /**
     * 存储向量数据
     */
    public void storeVector(KnowledgeDocumentSegment segment, List<Float> vector) {

        // 构建要插入的数据
        List<JsonObject> data = new ArrayList<>();
        JsonObject row = new JsonObject();
        row.addProperty(ID_FIELD, segment.getId());
        row.addProperty(DOCUMENT_ID_FIELD, segment.getDocumentId());
        row.addProperty(SEGMENT_ID_FIELD, segment.getId());
        row.addProperty(SEGMENT_INDEX_FIELD, (long) segment.getSegmentIndex());

        // 添加向量数据
        JsonArray vectorArray = new JsonArray();
        for (Float v : vector) {
            vectorArray.add(v);
        }
        row.add(VECTOR_FIELD, vectorArray);
        row.addProperty(CONTENT_FIELD, segment.getSegmentContent());

        // 构建元数据
        Map<String, Object> metadataMap = new HashMap<>();
        metadataMap.put("wordCount", segment.getWordCount());
        metadataMap.put("createTime", segment.getCreateTime());
        metadataMap.put("segmentTitle", segment.getSegmentTitle());
        row.addProperty(METADATA_FIELD, JSON.toJSONString(metadataMap));
        data.add(row);

        // 插入数据
        InsertResp insertResp = milvusClient.insert(InsertReq.builder()
                .collectionName(COLLECTION_NAME)
                .data(data)
                .build());
        log.info("插入向量数据库, 插入的数量: " + insertResp.getInsertCnt());
    }

    /**
     * 搜索相似向量
     */
    public List<SearchResult> searchSimilarVectors(List<Float> queryVector, int topK) {
        // 转换向量格式
        float[] floatArray = new float[queryVector.size()];
        for (int i = 0; i < queryVector.size(); i++) {
            floatArray[i] = queryVector.get(i);
        }

        // 创建搜索向量
        FloatVec floatVec = new FloatVec(floatArray);

        // 构造搜索请求
        SearchReq searchReq = SearchReq.builder()
                .collectionName(COLLECTION_NAME)
                .data(Collections.singletonList(floatVec))
                .annsField(VECTOR_FIELD)
                .topK(topK)
                .outputFields(Arrays.asList(ID_FIELD, DOCUMENT_ID_FIELD, SEGMENT_INDEX_FIELD, CONTENT_FIELD, METADATA_FIELD))
                .build();

        // 执行搜索
        SearchResp response = milvusClient.search(searchReq);

        // 处理搜索结果
        List<SearchResult> results = new ArrayList<>();
        List<List<io.milvus.v2.service.vector.response.SearchResp.SearchResult>> searchResults = response.getSearchResults();
        for (List<io.milvus.v2.service.vector.response.SearchResp.SearchResult> resultPerVector : searchResults) {
            for (io.milvus.v2.service.vector.response.SearchResp.SearchResult item : resultPerVector) {
                SearchResult result = new SearchResult();
                result.setId((String) item.getEntity().get(ID_FIELD));
                result.setDocumentId((String) item.getEntity().get(DOCUMENT_ID_FIELD));
                result.setSegmentIndex(((Long) item.getEntity().get(SEGMENT_INDEX_FIELD)).intValue());
                result.setContent((String) item.getEntity().get(CONTENT_FIELD));
                result.setScore(item.getScore());
                result.setMetadata((String) item.getEntity().get(METADATA_FIELD));
                results.add(result);
            }
        }

        return results;
    }


    /**
     * 删除指定文档的向量数据
     */
    public void deleteDocumentVectors(String documentId) {
        String filter = DOCUMENT_ID_FIELD + " == '" + documentId + "'";
        milvusClient.delete(DeleteReq.builder()
                .collectionName(COLLECTION_NAME)
                .filter(filter)
                .build());
        log.info("删除向量数据 通过 document ID: " + documentId);
    }

    /**
     * 根据多个文档ID删除向量数据
     */
    public void deleteDocumentVectorsByDocumentIds(List<String> documentIds) {
        if (documentIds == null || documentIds.isEmpty()) {
            log.warn("文档ID列表为空，无需删除操作");
            return;
        }

        // 构造过滤条件
        StringBuilder filterBuilder = new StringBuilder(DOCUMENT_ID_FIELD + " in [");
        for (int i = 0; i < documentIds.size(); i++) {
            filterBuilder.append("'").append(documentIds.get(i)).append("'");
            if (i < documentIds.size() - 1) {
                filterBuilder.append(",");
            }
        }
        filterBuilder.append("]");

        String filter = filterBuilder.toString();

        // 执行删除操作
        milvusClient.delete(DeleteReq.builder()
                .collectionName(COLLECTION_NAME)
                .filter(filter)
                .build());

        log.info("批量删除向量数据，文档IDs: " + documentIds);
    }


    /**
     * 搜索结果类
     */
    @Data
    public static class SearchResult {
        private String id;
        private String documentId;
        private int segmentIndex;
        private String content;
        private float score;
        private String metadata;
    }
}
